from solvers.solverbase import SolverBase
from solvers.analytical.analyticalbase import AnalyticalBase
from dolfin import *
from math import sqrt
import numpy

class GaussLegendre6Direct(SolverBase,AnalyticalBase):
    '3-symplectic method rewritten without increments.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        AnalyticalBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)
        phi3 = Function(V)        # solution to the aux. system  
        phi2 = Function(V)        # solution to the other aux. system 
        phi1 = Function(V)        # solution

        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        (ib,ib_value,bParts) = problem.get_inflowBCS()
        if problem.weakBC:
            raise NotImplementedError
        else:
            n = FacetNormal(mesh)
            a3 = phi*varphi*dx - dt/18.*3*sqrt(15.)/20.*(phi*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            a2 = phi*varphi*dx + dt/18.*3*sqrt(15.)/20.*(phi*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            a1 = phi*varphi*dx + dt/2.*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)

            l3 = phi_*varphi*dx + dt/18.*21.*sqrt(15.)/20.*(phi2*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi2*v)*dx - phi2*varphi*nabla_div(v)*dx)\
                                + dt/18.*(4.+3.*sqrt(15.)/5.)*(phi1*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi1*v)*dx - phi1*varphi*nabla_div(v)*dx)\
                                + dt/18.*(4.-9.*sqrt(15.)/5.)*(phi_*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)

            l2 = phi_*varphi*dx - dt/18.*(4.-3*sqrt(15.)/5.)*(phi_*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                                - dt/18.*(4.-3.*sqrt(15.)/5.)*(phi1*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi1*v)*dx - phi1*varphi*nabla_div(v)*dx)\
                                - dt/18.*(21.*sqrt(15.)/20.)*(phi3*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi3*v)*dx - phi3*varphi*nabla_div(v)*dx)

            l1 = phi_*varphi*dx - dt/18.*(9.-9*sqrt(15.)/5.)*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                                - dt/18.*(9.*sqrt(15.)/10.)*(phi2*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi2*v)*dx - phi2*varphi*nabla_div(v)*dx)\
                                - dt/18.*(9.*sqrt(15.)/10.)*(phi3*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi3*v)*dx - phi3*varphi*nabla_div(v)*dx)


        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        A3 = assemble(a3)
        A2 = assemble(a2)
        A1 = assemble(a1,exterior_facet_domains=bParts)
        
        phi = Function(V)
   
        t = float(dt)
        count = 1
       
        eps_ = -1.
        eps = 0.
        iteration = 0
        maxIter = 10
        tol = 1E-8

        phi2.assign(phi_) # initial guesses
        phi1.assign(phi_)
        while t <= T:
            ibc = DirichletBC(V,ib_value(t),ib)
            while abs(eps_-eps) > tol and iteration < maxIter:
                iteration += 1
                L3 = assemble(l3)
                solve(A3,phi.vector(),L3)  # solve system and assign the solution to phi3
                phi3.assign(phi)

                L2 = assemble(l2)
                solve(A2,phi.vector(),L2)  # solve system and assign the solution to phi2
                phi2.assign(phi)

                L1 = assemble(l1,exterior_facet_domains=bParts)
                ibc.apply(A1,L1)
                solve(A1,phi.vector(),L1)  # solve system and assign the solution to phi1
                phi1.assign(phi)

                diff = phi3.vector().array() - phi2.vector().array() - phi1.vector().array()
                eps_ = eps
                eps = numpy.linalg.norm(diff,ord=2)        # compute the error

            phi_.assign(phi1)
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1
            eps_ = -1.
            eps = 0.
            iteration = 0

        problem.save_data(self.saveDir,str(self.CFL))
        
#<-------------------------------------------TRICKY THINGS----------------------------------------------------s-
        
class GaussLegendre6(SolverBase,AnalyticalBase):
    '3-symplectic method as found on wikipedia. Fails with boundary condiotions. Avoid!'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        AnalyticalBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)

        k1 = Function(V)                   
        k2 = Function(V)
        k3 = Function(V)

        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        (ib,ib_value,ibParts) = problem.get_inflowBCS()
        if problem.weakBC:
            a1 = phi*varphi*dx - 5./36.*dt*(inner(nabla_grad(varphi),phi*v)*dx + phi*varphi*nabla_div(v)*dx)
            a2 = phi*varphi*dx - 2./9.*dt*(inner(nabla_grad(varphi),phi*v)*dx + phi*varphi*nabla_div(v)*dx)
            a3 = phi*varphi*dx - 5./36.*dt*(inner(nabla_grad(varphi),phi*v)*dx + phi*varphi*nabla_div(v)*dx)
           
            l1 =                        dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                  + dt*(2./9.-sqrt(15.)/15.)*(inner(nabla_grad(varphi),k2*v)*dx + k2*varphi*nabla_div(v)*dx)\
                  + dt*(5./36.+sqrt(15.)/30.)*(inner(nabla_grad(varphi),k3*v)*dx + k3*varphi*nabla_div(v)*dx)
       
       
            l2 =                        dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                  + dt*(5./36.+sqrt(15.)/24.)*(inner(nabla_grad(varphi),k1*v)*dx + k1*varphi*nabla_div(v)*dx)\
                  + dt*(5./36.-sqrt(15.)/24.)*(inner(nabla_grad(varphi),k3*v)*dx + k3*varphi*nabla_div(v)*dx)      
       
            l3 =                        dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                  + dt*(5./36.+sqrt(15.)/30.)*(inner(nabla_grad(varphi),k1*v)*dx + k1*varphi*nabla_div(v)*dx)\
                  + dt*(2./9.+sqrt(15.)/15.)*(inner(nabla_grad(varphi),k2*v)*dx + k2*varphi*nabla_div(v)*dx)     
        else:
            n = FacetNormal(mesh)
            a1 = phi*varphi*dx + 5./36.*dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            a2 = phi*varphi*dx + 2./9.*dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            a3 = phi*varphi*dx + 5./36.*dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
           
            l1 =  -                      dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                  - dt*(2./9.-sqrt(15.)/15.)*(k2*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k2*v)*dx - k2*varphi*nabla_div(v)*dx)\
                  - dt*(5./36.+sqrt(15.)/30.)*(k3*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k3*v)*dx - k3*varphi*nabla_div(v)*dx)
       
       
            l2 =  -                      dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                  - dt*(5./36.+sqrt(15.)/24.)*(k1*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k1*v)*dx - k1*varphi*nabla_div(v)*dx)\
                  - dt*(5./36.-sqrt(15.)/24.)*(k3*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k3*v)*dx - k3*varphi*nabla_div(v)*dx)      
       
            l3 =  -                      dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                  - dt*(5./36.+sqrt(15.)/30.)*(k1*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k1*v)*dx - k1*varphi*nabla_div(v)*dx)\
                  - dt*(2./9.+sqrt(15.)/15.)*(k2*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k2*v)*dx - k2*varphi*nabla_div(v)*dx)     
       
        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        k2.assign(interpolate(Constant(0),V)) 
        k3.assign(k2)

        A1 = assemble(a1,exterior_facet_domains=ibParts)
        A2 = assemble(a2,exterior_facet_domains=ibParts)
        A3 = assemble(a3,exterior_facet_domains=ibParts)

        phi = Function(V)
   
        t = float(dt)
        count = 1
        
        iteration = 0
        maxIter = 50
        tol = 1E-8
        eps_ = -1.
        eps = 0.

        while t <= T:
            ibc = DirichletBC(V,Constant(0.),ib)
            while abs(eps_-eps) > tol and iteration < maxIter:
                iteration += 1
                L1 = assemble(l1,exterior_facet_domains=ibParts)
                ibc.apply(A1,L1)
                solve(A1,phi.vector(),L1)
                k1.assign(phi)

                L2 = assemble(l2,exterior_facet_domains=ibParts)
                ibc.apply(A2,L2)
                solve(A2,phi.vector(),L2)
                k2.assign(phi)

                L3 = assemble(l3,exterior_facet_domains=ibParts)
                ibc.apply(A3,L3)
                solve(A3,phi.vector(),L3)
                k3.assign(phi)

                diff = k1.vector().array() - k2.vector().array() - k3.vector().array() 
                
                eps_ = eps
                eps = numpy.linalg.norm(diff,ord=2)    
           
            ibc = DirichletBC(V,ib_value(t),ib)
            ibc.apply(phi_.vector())
            phi_.vector()[:] = phi_.vector().array() + (5./18.)*k1.vector().array() + (8./18.)*k2.vector().array() + (5./18.)*k3.vector().array()
            
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1
            eps_ = -1.
            eps = 0
            iteration = 0

        problem.save_data(self.saveDir,str(self.CFL))


